import numpy as np
from skimage.io import imread
import matplotlib.pyplot as plt
import pickle
from sklearn.metrics import accuracy_score

# 加载图像
X_img = imread(r"C:\Users\29650\Desktop\yanashi\Sandstone_2.tif")
y_true = imread(r"C:\Users\29650\Desktop\yanashi\Sandstone_2_segment.tif")

print("X_img shape:", X_img.shape)
print("y_true shape:", y_true.shape)

# 如果y_true是三通道的，我们只取第一个通道
if len(y_true.shape) == 3:
    y_true = y_true[:,:,0]  # 只取第一个通道

# 确保两个图像具有相同的尺寸
min_height = min(X_img.shape[0], y_true.shape[0])
min_width = min(X_img.shape[1], y_true.shape[1])

X_img = X_img[:min_height, :min_width]
y_true = y_true[:min_height, :min_width]

# 将图像重塑为二维数组
X = X_img.reshape(X_img.shape[0] * X_img.shape[1], -1)

# 加载训练好的模型
with open('clf.pkl', 'rb') as f:
    clf = pickle.load(f)

# 预测分割
y_pred = clf.predict(X)

# 计算准确率
acc = accuracy_score(y_true.flatten(), y_pred)
print(f"分割准确率: {acc:.3f}")

# 显示结果
fig, axes = plt.subplots(1, 3, figsize=(15, 5))

axes[0].imshow(X_img, cmap='gray')
axes[0].set_title('original image')

axes[1].imshow(y_true, cmap='gray')
axes[1].set_title('segement')

axes[2].imshow(y_pred.reshape(y_true.shape), cmap='gray')
axes[2].set_title(f'segementation (acc: {acc:.3f})')

plt.show()